from general import *
from search.utils import PriorityQueue
'''
Created on Apr 21, 2011

@author: bob
'''

class Pile_Division(object):
    '''
    classdocs
    '''

    def __init__(self, pState):
        '''
        Constructor
        '''
        self.obstacles = {}
        for o in pState.obstacle_locations:
            self.obstacles[o] = 1
            
        piles = {}
        self.piles_to_robots = {}
        for p in pState.dirt_locations:
            piles[p] = 1
        
#used to divide piles between robots        
        self.robots_piles = {} 
        robots_prios = {}
        for i in range(0,len(pState.robots)):
            def evaluator(pile):
                return self.evalPaths(pState.robots[i],pile)
            q = PriorityQueue(evaluator)
            for p in pState.dirt_locations:
                q.append(p)
            robots_prios[i] = q
            self.robots_piles[i] = set()
        
        while (len(piles) != 0):
            for i in range(0,len(pState.robots)):
                #break when no more piles
                if (len(piles) == 0):
                    break;
                #get next best pile for current robot                
                while(True):
                    if (len(robots_prios[i])==0):
                        break
                    p = robots_prios[i].pop()
                    if p in piles:
                        self.robots_piles[i].add(p)
                        self.piles_to_robots[p] = i
                        del piles[p]
                        break
        
#        print self.robots_piles        
    
    
    def tryAssignPile(self, r, pState):
        if len(pState.dirt_locations)==0:
            return None
        
#       find the pile in max dist from its robot
        maxDist = -infinity
        maxPile = None 
        for p in pState.dirt_locations:
            temp = self.evalPaths(
                pState.robots[self.piles_to_robots[p]], p)
            if temp > maxDist:
                maxPile = p
                maxDist = temp
#       take the pile from its master only if the master has more piles
#       or it is better to pass the pile
        if (len(self.robots_piles[self.piles_to_robots[maxPile]])==1 and
            self.evalPaths(pState.robots[r], maxPile) >= maxDist):
            return None
#        print maxPile,self.piles_to_robots[maxPile],"to"+str(r)
        self.robots_piles[self.piles_to_robots[maxPile]].discard(p)
        self.robots_piles[r].add(maxPile)
        self.piles_to_robots[maxPile] = r 
        return  (maxPile, r)
#   count obstacles along manhattan distances     
    def evalPaths(self,r, p):
        if (r[0]>p[0]):
            xStart,xEnd, yStart, yEnd, yStep = p[0],r[0], p[1], r[1], 1
            if (p[1]>r[1]):
                yStep = -1
                yEnd += -2
        else:
            xStart,xEnd, yStart, yEnd, yStep = r[0],p[0], r[1], p[1], 1
            if (r[1]>p[1]):
                yStep = -1
                yEnd += -2
#       path 1:
        cost1 = manhattan(r, p)
        for y in range(yStart,yEnd+1, yStep):
            if ((xStart, y) in self.obstacles):
                cost1 += 1
        for x in range(xStart,xEnd+1):
            if ((x, yEnd) in self.obstacles):
                cost1 += 1
#           path 2: 
        cost2 = manhattan(r, p)
        for x in range(xStart,xEnd+1):
            if ((x, yStart) in self.obstacles):
                cost2 += 1
        for y in range(yStart,yEnd+1, yStep):
            if ((xEnd, y) in self.obstacles):
                cost2 += 1         
        return min(cost1, cost2)
    
    def evaluate(self,pState):
        totalCost = 0
        for i in range(0,len(pState.robots)):
            skipped = 0
            roboCost = 0
#            if (len(self.robots_piles[i])==0):
#                self.tryAssignPile(i, pState)
            for p in self.robots_piles[i]:
                if not p in pState.dirt_locations:
                    skipped += 1
                    continue
#                roboCost += self.evalPaths(pState.robots[i], p)
                roboCost += self.evalPaths(pState.robots[i], p)
#            print self.robots_piles[i]
            totalCost += roboCost*(len(self.robots_piles[i])-skipped)
#        totalCost = 0
#        for p in pState.dirt_locations:
#            totalCost += manhattan(
#                pState.robots[self.piles_to_robots[p]], p)
#        if totalCost != totalCost2:
#            print totalCost,totalCost2
        return totalCost
            